import os, pickle, sys
import soundfile as sf
from SFR_Base import SFR_Base
import numpy as np

Help_info = '''
#----------------------------------声场重构核心程序-----------------------------------------
0.帮助文档:    显示当前帮助文档

-H -help
eg: 
.\SFR_main.exe -H


1.全局计算：
1.1 不带标识符的默认参数,直接进行全局计算。计算结果将保存到默认路径下，涉及结果对象：sfr.pkl 和 计算得到的扬声器输出信号（用于重构的最终信号）：Speaker_out.wav

[rand_in_path] [response_path] [target_path] [save_obj_path][save_wav_path][filter_len] [beta] [valid_filter_ratio]
eg1: 
.\SFR_main.exe C:/Users/TXH/Desktop/SFR/wavPinkNoise_24bit_48k_20s.wav C:/Users/TXH/Desktop/SFR/Channel_64.wav E:/Pywork/Sound_reconstruction/data/Airport_Handset.wav C:/Users/TXH/Desktop/sfr.pkl C:/Users/TXH/Desktop/Speaker_out.wav 14400 3e-2 0.5

1.2 只计算全局计算中的第一部分，获取逆传递函数。
[rand_in_path] [response_path][save_obj_path][filter_len] [beta] [valid_filter_ratio]
eg2: 
.\SFR_main.exe C:/Users/TXH/Desktop/SFR/wavPinkNoise_24bit_48k_20s.wav C:/Users/TXH/Desktop/SFR/Channel_64.wav C:/Users/TXH/Desktop/sfr.pkl 14400 3e-2 0.5


2.调试： 读取本地保存的历史计算结果 sfr.pkl 对象文件，调整滤波器长度和正则化因子等参数，重新计算输出结果，以满足条件。

-D [beta][valid_filter_ratio][load_obj_path][target_path][save_wav_path]
eg: 
.\SFR_main.exe -D 4e-3 0.5 C:/Users/TXH/Desktop/sfr.pkl E:/Pywork/Sound_reconstruction/data/Airport_Handset.wav C:/Users/TXH/Desktop/Speaker_out.wav

注意：将各路径参数更改为实际绝对路径。


3.独立计算功能:
3.1 计算1/3倍频程，获得中心频率和频带能量数据,结果保存为.data后缀的文件。

-F_PCB [wav_path] [sample_rate] [save_data_name]     
eg: 
.\SFR_main.exe -F_PCB C:/Users/TXH/Desktop/SFR/random_in.wav 48000 C:/Users/TXH/Desktop/random_in_CPB.data 


3.2 读取默认路径中储存的 C:/Users/TXH/Desktop/sfr.pkl 文件，逆滤波得到输出信号,输出结果保存到指定的save_wav文件中。

-F_Inv_filtering [target_wav_path] [load_obj_path] [save_wav_path] 
eg: 
.\SFR_main.exe -F_Inv_filtering E:/Pywork/Sound_reconstruction/data/Airport_Handset.wav C:/Users/TXH/Desktop/sfr.pkl C:/Users/TXH/Desktop/Airport_Handset_Speaker_out.wav


'''
# 使用方法示例：
# argv_list = ['SFR_main.exe','-H']
# argv_list = ['SFR_main.exe','-F_PCB','C:/Users/TXH/Desktop/SFR/wavPinkNoise_24bit_48k_20s.wav',48000,'C:/Users/TXH/Desktop/random_in_CPB.data']
argv_list = ['SFR_main.exe', 'wavPinkNoise_24bit_48k_20s.wav',
             'Airport_Handset.wav',
             'Speaker_out.wav', 14400, 1e-5, 0.5]
# argv_list = ['SFR_main.exe','C:/Users/TXH/Desktop/SFR/wavPinkNoise_24bit_48k_20s.wav','C:/Users/TXH/Desktop/SFR/Channel_64.wav','C:/Users/TXH/Desktop/sfr.pkl',14400,3e-2,0.5]
# argv_list = ['SFR_main.exe','-D',4e-3,0.5,'C:/Users/TXH/Desktop/sfr.pkl','E:/Pywork/Sound_reconstruction/data/Airport_Handset.wav','C:/Users/TXH/Desktop/Speaker_out.wav']
# argv_list = ['SFR_main.exe','-F_Inv_filtering','E:/Pywork/Sound_reconstruction/data/Airport_Handset.wav','C:/Users/TXH/Desktop/sfr.pkl','C:/Users/TXH/Desktop/Airport_Handset_Speaker_out.wav']

# ----------------------------------------------------------------------------------------
try:
    os.mkdir('./temp')
except:
    pass

# argv_list = sys.argv
sfr = SFR_Base()

if argv_list[1] == '-H' or argv_list[1] == '-help':
    print(Help_info)
elif argv_list[1] == '-F_PCB':
    if len(argv_list) == 5:
        freq, out = sfr.CPB(sf.read(argv_list[2])[0], int(argv_list[3]))
        sfr.save_data(argv_list[4], np.c_[freq, out])
    else:
        raise 'number of inputs should be 5!'

elif argv_list[1] == '-F_Inv_filtering':
    if len(argv_list) == 5:
        try:
            with open(argv_list[3], 'rb') as f:
                sfr3 = pickle.load(f)
        except:
            raise 'Can not find local file \'sfr.pkl\' in temporal dir!'

        target = sf.read(argv_list[2])[0]
        Speaker_out, amp_max = sfr.Inv_filtering(target, sfr3.c_list)
        print('**Max output amplitude: {:.2f}'.format(amp_max))
        sfr.save_data(argv_list[4], Speaker_out.T)
    else:
        raise 'number of inputs should be 5!'

elif argv_list[1] == '-D':
    if len(argv_list) == 7:
        try:
            with open(argv_list[4], 'rb') as f:
                sfr2 = pickle.load(f)
        except:
            raise 'Can not find local file \'sfr.pkl\' in temporal dir!'

        target = sf.read(argv_list[5])[0]
        # 计算逆传递函数和扬声器的输出
        sfr2.c_list, C_list = sfr.Cal_inversion(sfr2.H_list, fs=sfr2.fs, pre_noise_len=sfr2.pre_noise_len,
                                                beta=float(argv_list[2]),
                                                keep_filter=float(argv_list[3]))
        Speaker_out, amp_max = sfr.Inv_filtering(target, sfr2.c_list)
        sfr.save_data(argv_list[6], Speaker_out.T)
        # 仿真检验
        sfr2.Error = sfr.Validate(target, sfr.Cal_sim_P(Speaker_out, sfr2.h_list))
        print('**Max Speaker Output: {:.2f}'.format(amp_max))  # 打印最大值
        print('**SPL Error:%.2f dB' % (sfr2.Error.mean()))

        # 保存中间量
        file_name, ext = os.path.splitext(argv_list[4])
        with open(file_name + '_new' + ext, 'wb') as f:
            pickle.dump(sfr2, f)  # 保存含有传递函数、滤波器长度、采样率等结果的对象
    else:
        raise 'number of inputs should be 7!'
else:
    if len(argv_list) == 9:
        rand_in_path, response_path, target_path, save_obj_path, save_wav_path, filter_len, beta, keep_filter = argv_list[
                                                                                                                1:]
    elif len(argv_list) == 7:
        rand_in_path, response_path, save_obj_path, filter_len, beta, keep_filter = argv_list[1:]
        target_path, save_wav_path = '', ''
    else:
        raise 'number of inputs should be 9 or 6 in default mode! Read the help doc.'

    # 初始化含数据实例，并读取传递函数建模数据
    beta, beta2 = float(beta), 0  # 初始化 beta 值
    iter, max_iter = 0, 100  # 最大迭代次数设置
    amp_error, min_Error = 1, 1e-3  # 初始化误差及误差阈值
    flag = True
    log_data = ['iter,beta,max_amp,Error_in_dB\n']

    rand_in, response, target = sfr.read_wav_data(rand_in_path=rand_in_path, response_path=response_path,
                                                  target_path=target_path, filter_len=int(filter_len))
    sfr.H_list, sfr.h_list = sfr.Cal_Impulse(rand_in[:96000], response[:, :, :96000], sfr.fs, sfr.filter_len,
                                             sfr.pre_noise_len, keep_filter=float(keep_filter))
    sfr.c_list, C_list = sfr.Cal_inversion(sfr.H_list, fs=sfr.fs, pre_noise_len=sfr.pre_noise_len, beta=beta,
                                           keep_filter=float(keep_filter))

    if len(argv_list) == 9:
        # 初始化逆滤波
        Speaker_out, amp_max = sfr.Inv_filtering(target, sfr.c_list)  # 逆滤波

        ## 粗调，确定一个数量级的空间
        while flag:
            iter += 1
            beta2 = beta * 10 if amp_max > 1 else beta / 10
            sfr.c_list, C_list = sfr.Cal_inversion(sfr.H_list, fs=sfr.fs, pre_noise_len=sfr.pre_noise_len, beta=beta2,
                                                   keep_filter=float(keep_filter))  # 计算逆矩阵
            Speaker_out, amp_max2 = sfr.Inv_filtering(target, sfr.c_list)  # 逆滤波
            # Level = np.max([np.abs(hh * (cc @ pp)) for hh, cc, pp in zip(sfr.H_list, C_list, F_target)])
            # print('Cycle {:d}: beta:{:.4e}  Max_output:{:.4f} max_level:{:.4f}'.format(iter, beta2, amp_max2,Level))  # 打印最大值
            sfr.Error = sfr.Validate(target, sfr.Cal_sim_P(Speaker_out, sfr.h_list))  # 仿真检验
            sfr.save_data(save_wav_path, Speaker_out.T)  # 保存计算的X输出
            print('**SPL Error:%.4f dB' % (sfr.Error.mean()))
            log_data.append(str(iter) + ',' + str(beta2) + ',' + str(amp_max2) + ',' + str(sfr.Error.mean()) + '\n')
            sfr.write_log_data(log_data)

            if (amp_max - 1) * (amp_max2 - 1) < 0:
                flag = False
                break
            else:
                amp_max, beta = amp_max2, beta2

        ## 二分法细调
        while iter < max_iter and amp_error > min_Error:
            iter += 1
            beta3 = (beta + beta2) / 2
            sfr.c_list, C_list = sfr.Cal_inversion(sfr.H_list, fs=sfr.fs, pre_noise_len=sfr.pre_noise_len, beta=beta3,
                                                   keep_filter=float(keep_filter))
            Speaker_out, amp_max3 = sfr.Inv_filtering(target, sfr.c_list)  # 逆滤波
            # Level = np.max([np.abs(hh * (cc @ pp)) for hh, cc, pp in zip(sfr.H_list, C_list, F_target)])
            amp_error = abs(amp_max3 - 1)
            # print('Cycle {:d}: beta:{:.4e}  Max_output:{:.4f} max_level:{:.4f}'.format(iter, beta3, amp_max3,Level))
            sfr.Error = sfr.Validate(target, sfr.Cal_sim_P(Speaker_out, sfr.h_list))  # 仿真检验
            sfr.save_data(save_wav_path, Speaker_out.T * 0.5)  # 保存计算的X输出
            print('**SPL Error:%.4f dB' % (sfr.Error.mean()))

            beta2 = max(beta, beta2) if amp_max3 > 1 else min(beta, beta2)  # 保留 beta3 和上/下限值
            beta = beta3
            log_data.append(str(iter) + ',' + str(beta3) + ',' + str(amp_max3) + ',' + str(sfr.Error.mean()) + '\n')
            sfr.write_log_data(log_data)
        log_data.append('Done!')
        sfr.write_log_data(log_data)

    # 保存中间量
    with open(save_obj_path, 'wb') as f:
        pickle.dump(sfr, f)  # 保存含有传递函数、滤波器长度、采样率等结果的对象

print('done!')

# import matplotlib.pyplot as plt
# for i in range(8):
#     # plt.plot(target[:,i])
#     plt.plot(Speaker_out[i, :])
# plt.show()
